package com.xjgz.aio.socket.pro.extension.plugins;

import com.xjgz.aio.socket.pro.channels.AsynchronousSocketChannelProxy;

import java.nio.ByteBuffer;
import java.nio.channels.AsynchronousSocketChannel;
import java.nio.channels.CompletionHandler;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

public class RateLimiterPlugin<T> extends AbstractPlugin<T> {

    /**
     * read 流控阈值
     */
    private final int readRateLimiter;

    /**
     * write 流控阈值
     */
    private final int writeRateLimiter;

    /**
     * 是否启用流控
     */
    private final boolean enabled;

    private final int bufferTime = 10;
    private ScheduledExecutorService executorService;

    public RateLimiterPlugin(int readRateLimiter, int writeRateLimiter) {
        this.readRateLimiter = readRateLimiter;
        this.writeRateLimiter = writeRateLimiter;
        this.enabled = readRateLimiter > 0 && writeRateLimiter > 0;
        if (enabled) {
            executorService = Executors.newSingleThreadScheduledExecutor();
        }
    }

    @Override
    public AsynchronousSocketChannel shouldAccept(AsynchronousSocketChannel channel) {
        return enabled ? new RateLimiterChannel(channel, readRateLimiter, writeRateLimiter) : channel;
    }

    class RateLimiterChannel extends AsynchronousSocketChannelProxy {
        private final int readRateLimiter;
        private final int writeRateLimiter;

        /**
         * 上一次read流控窗口临界点
         */
        private long latestReadTime;

        /**
         * 流控窗口期输入字节数
         */
        private int readCount;

        /**
         * 上一次write流控窗口临界点
         */
        private long latestWriteTime;

        /**
         * 流控窗口期输出字节数
         */
        private int writeCount;

        public RateLimiterChannel(AsynchronousSocketChannel asynchronousSocketChannel, int readRateLimiter, int writeRateLimiter) {
            super(asynchronousSocketChannel);
            this.readRateLimiter = readRateLimiter;
            this.writeRateLimiter = writeRateLimiter;
        }

        @Override
        public <A> void read(ByteBuffer dst, long timeout, TimeUnit unit, A attachment, CompletionHandler<Integer, ? super A> handler) {
            if (dst.remaining() == 0 || readRateLimiter <= 0) {
                super.read(dst, timeout, unit, attachment, handler);
                return;
            }
            int availableSize;
            long remainingTime = 1000 + latestReadTime - System.currentTimeMillis();
            // 新的流控窗口
            if (remainingTime <= bufferTime) {
                readCount = 0;
                latestReadTime = System.currentTimeMillis();
            }
            availableSize = Math.min(readRateLimiter - readCount, dst.remaining());
            // 触发流控
            if (availableSize <= 0) {
                executorService.schedule(() -> RateLimiterChannel.this.read(dst, timeout, unit, attachment, handler), remainingTime, TimeUnit.MILLISECONDS);
                return;
            }

            int limit = dst.limit();
            dst.limit(dst.position() + availableSize);
            super.read(dst, timeout, unit, attachment, new CompletionHandler<>() {
                @Override
                public void completed(Integer result, A attachment) {
                    if (result > 0) {
                        if (System.currentTimeMillis() - latestReadTime > 1000) {
                            readCount = 0;
                            latestReadTime = System.currentTimeMillis();
                        } else {
                            readCount += result;
                        }
                    }
                    dst.limit(limit);
                    handler.completed(result, attachment);
                }

                @Override
                public void failed(Throwable exc, A attachment) {
                    handler.failed(exc, attachment);
                }
            });
        }

        @Override
        public <A> void write(ByteBuffer src, long timeout, TimeUnit unit, A attachment, CompletionHandler<Integer, ? super A> handler) {
            if (src.remaining() == 0 || writeRateLimiter <= 0) {
                super.write(src, timeout, unit, attachment, handler);
                return;
            }

            int availableWriteSize;
            long remainingTime = 1000 + latestWriteTime - System.currentTimeMillis();
            if (remainingTime <= bufferTime) {
                writeCount = 0;
                latestWriteTime = System.currentTimeMillis();
            }
            availableWriteSize = Math.min(writeRateLimiter - writeCount, src.remaining());
            if (availableWriteSize <= 0) {
                executorService.schedule(() -> RateLimiterChannel.this.write(src, timeout, unit, attachment, handler), remainingTime, TimeUnit.MILLISECONDS);
                return;
            }

            int limit = src.limit();
            src.limit(src.position() + availableWriteSize);
            super.write(src, timeout, unit, attachment, new CompletionHandler<>() {
                @Override
                public void completed(Integer result, A attachment) {
                    if (result > 0) {
                        if (System.currentTimeMillis() - latestWriteTime > 1000) {
                            writeCount = 0;
                            latestWriteTime = System.currentTimeMillis();
                        } else {
                            writeCount += result;
                        }
                    }
                    src.limit(limit);
                    handler.completed(result, attachment);
                }

                @Override
                public void failed(Throwable exc, A attachment) {
                    handler.failed(exc, attachment);
                }
            });
        }
    }
}
